#!/usr/bin/python
import sys, os, StringIO, imp
from struct import *
from optparse import *

def getCompressors(exedir):
	files = os.listdir(exedir)
	compressorFileList = []
	compressorDict = {}
	for file in files:
		if file.endswith("Compressor.py") or file.endswith("Compressor.pyc"):
			compressorFileList.append(os.path.splitext(file)[0])
	for module in compressorFileList:
		tempModule = __import__(module)
		compressorClass = tempModule.__dict__[module]
		compressorDict[module.lower()] = compressorClass
	return compressorDict

class DefaultChunker:
	def __init__(self, chunksize):
		self.streamObj = None
		self.end = None
		self.chunksize = chunksize
		self.done = True
		
	def process(self, streamObj):
		self.streamObj = streamObj
		self.done = False
		
	def hasChunks(self):
		return not self.done

		
	def nextChunk(self):
		text = self.streamObj.read(self.chunksize)
		if len(text) == 0:
			self.done = True
		return (text, len(text))

class ZipControl:
	MAJOR = 0
	MINOR = 1
		
	def __createZipChunk(self, unzippedData, unzippedDataLen, writer, headerList):
		zippedData = self.zipper.compress(unzippedData)
		zippedDataLen = len(zippedData)
		writer.write(zippedData)
		headerList.append((unzippedDataLen, zippedDataLen))
		
	def __createHeader(self, thelist):
		zipClass = str(self.compressorName)
		zipClassLen = len(zipClass)
		packString = 'bbbbbb%ds' %zipClassLen
		header = pack(packString,0x52,0x41,0x5a,
		                self.MAJOR, self.MINOR, 
				zipClassLen, zipClass)
		header += pack('l',len(thelist))
		for (unzipsize, zipsize) in thelist:
			header += pack('ll',unzipsize, zipsize)
		return header
	
	def __init__(self, compressorName, chunker, zipChunkSize=10000000):
		self.chunker = chunker
		self.compressorName = compressorName
		self.zipper = g_Compressors[compressorName]()
		self.maxZipChunkSize = zipChunkSize
		
	def compress(self, streamObj, outputFilename):
		zipped = []
		noheaderFilename = os.path.join("/tmp/",outputFilename+".noheader")
		writer = open(noheaderFilename,"wb+")
		self.chunker.process(streamObj)
		
		zipableChunk = ""
		zipableChunkLen = 0
		killswitch=0
		while self.chunker.hasChunks():
			(chunk, chunkLen) = self.chunker.nextChunk()
			if zipableChunkLen > 0 and zipableChunkLen + chunkLen > self.maxZipChunkSize:
				self.__createZipChunk(zipableChunk, zipableChunkLen, writer, zipped)
				zipableChunk = chunk
				zipableChunkLen = chunkLen
			else:
				zipableChunk += chunk
				zipableChunkLen += chunkLen
			print zipableChunkLen, self.maxZipChunkSize
		if zipableChunkLen > 0:
			self.__createZipChunk(zipableChunk, zipableChunkLen, writer, zipped)
		writer.close()
		header = self.__createHeader(zipped)
		realWriter = open(outputFilename,"wb+")
		noheader = open(noheaderFilename)
		data = header
		while len(data) > 0:
			realWriter.write(data)
			data = noheader.read(4096)
		
	## static method ##
	def extractHeader(razfile):
		zipreader = open(razfile)
		try:
			(c0, c1, c2, major, minor, compressorStrSize) = unpack('bbbbbb', zipreader.read(6))
		except Exception,e:
			raise Exception("Incorrect Header. Unpack error: %s" % str(e))
		if c0 != 0x52 and c1 != 0x41 and c2 != 0x5a:
			raise Exception("Incorrect Header. Expected 0x52,0x41,0x5a identifier")
		try:
			(compressorStr,) = unpack('%ds'%compressorStrSize, zipreader.read(compressorStrSize))
		except Exception,e:
			raise Exception("Incorrect Header. Could not extract compressor class name")
		try:
			(listLen,) = unpack('l',zipreader.read(4))
		except Exception,e:
			raise Exception("Incorrect Header. Could not extract pointer list length")
		pointerList = []
		for counter in range(0, listLen):
			try:
				pointerList.append(unpack('ll',zipreader.read(8)))
			except Exception,e:
				raise Exception("Incorrect Heder. Could not extract pointer")
		location = zipreader.tell()
		zipreader.close()
		return (major, minor, compressorStr, pointerList, location)
	extractHeader = staticmethod(extractHeader)
			
	def decompress(self, razipfile, outputFilename):
		(major, minor, compressorStr, pointerList, start) = self.extractHeader(razipfile)
		"""In the future, check version"""
		if not g_Compressors.has_key(compressorStr):
			raise Exception("Cannot unzip %s. Requires %s." % (razipfile, compressorStr))
		decompressor = g_Compressors[compressorStr]()
		unzipwriter = open(outputFilename,"w+")
		zipreader = open(razipfile,"rb")
		zipreader.seek(start)
		for (unzipsize, zipsize) in pointerList:
			zippedData = zipreader.read(zipsize)
			unzipwriter.write(decompressor.decompress(zippedData))
		unzipwriter.close()
		zipreader.close()
		
	## static method ##
	def open(razipfileName):
		return RandomAccessZipFile(razipfileName)
	open = staticmethod(open)
			
"""
This is a read-only random-accessor for a RandomAccessZipFile

Follows the "file" interface
"""
class RandomAccessZipFile:
	class ChunkCache:
		def __init__(self):
			self.chunk = None
			self.begin = None
			self.end = None
			
		def cache(self, chunk, begin, end):
			self.chunk = chunk
			self.begin = begin
			self.end = end
		
		def seekGlobalPosition(self, position):
			if self.chunk == None: return None
			if position < self.begin or position > self.end: return None
			self.chunk.seek(position-self.begin)
			return self.chunk
			
	def __getZipBlockReader(self, globalPosition):
		if globalPosition < 0 or globalPosition > self.__end:
			raise Exception("seek out of bounds!")
		
		cached = self.__cachedChunk.seekGlobalPosition(globalPosition)
		if cached != None:
			return cached
		
		blockstart = 0
		unzipPointer = 0
		curBlockSize = 0
		end = None
		for (unzipsize, zipsize) in self.__pointerList:
			curBlockSize = zipsize
			if unzipsize + unzipPointer > globalPosition:
				end = unzipsize + unzipPointer
				break
			blockstart += zipsize
			unzipPointer += unzipsize
		if end == None:
			end == self.__end
		self.__realfile.seek(self.__start+blockstart)
		zippedData = self.__realfile.read(curBlockSize)
		unzippedData = self.__compressor.decompress(zippedData)
		reader = StringIO.StringIO(unzippedData)
		reader.seek(globalPosition-unzipPointer)
		
		self.__cachedChunk.cache(reader, blockstart, end)
		return reader
	
	def __init__(self, filename):
		(major, minor, compressorStr, self.__pointerList, self.__start) = \
				ZipControl.extractHeader(filename)
		if not g_Compressors.has_key(compressorStr):
			raise Exception("Cannot unzip %s. Requires %s." % (razipfile, compressorStr))
		self.__realfile = open(filename, "rb")
		self.__cachedChunk = RandomAccessZipFile.ChunkCache()
		
		self.__compressor = g_Compressors[compressorStr]()
		self.__closed = False
		self.__pos = 0
		self.__end = 0
		for (unzipsize, zipsize) in self.__pointerList:
			self.__end += unzipsize

	def close(self):
		self.__closed = True
		
	def flush(self):
		pass
	
	"""
	def fileno(self):
		
	def isatty(self):
		
		These methods are NOT provided as per 
		http://docs.python.org/lib/bltin-file-objects.html
	"""
	
	def next(self):
		pass
	
	"""
	With no arguments, this method will read to the end of a 
	zipBlock at the current file pointer position. If an explicit
	size is specified, it will read that many bytes across zipBlocks
	"""
	def read(self, maxSize = -1):
		zipBlockReader = self.__getZipBlockReader(self.__pos)
		data = zipBlockReader.read(maxSize)
		self.__pos += len(data)
		while maxSize >= 0 and len(data) < maxSize:
			zipBlockReader = self.__getZipBlockReader(self.__pos)
			nextData = zipBlockReader.read(maxSize - len(data))
			self.__pos += len(nextData)
			data += nextData
		return data
		
	def seek(self, offset, whence = 0):
		if whence == 0:
			self.__pos = offset
		elif whence == 1:
			self.__pos += offset
		else:
			self.__pos = self.__end + offset
		if self.__pos < 0 or self.__pos > self.__end:
			raise Exception("Invalid seek.")
		
		
			
if __name__ == "__main__":
	g_Compressors = getCompressors(os.path.dirname(sys.argv[0]))
	
	parser = OptionParser()
	parser.add_option("-m", "--mode", dest="mode",
			help="[(c)ompress/(d)ecompress]")
	parser.add_option("-s", "--stdin", dest="stdin",
			action="store_true", default=False,
			help="use stdin")
	#parser.add_option("-s", "--stdout", dest="stdout",
	#		action="store_false", default=False,
	#		help="use stdout")
	parser.add_option("-c", "--compressor", dest="compressor",
			help="Zip compressor to use",
			default="zlib")
	parser.add_option("-k", "--chunker", dest="chunker",
			help="later",
			default=None,
			metavar="FILE")
	parser.add_option("-i","--in",dest="infile",
			help="input",
			default=None,
			metavar="FILE")
	parser.add_option("-o","--out",dest="outfile",
			help="output",
			default=None,
			metavar="FILE")
	(options, args) = parser.parse_args()

	if options.chunker != None:
		moduleName = os.path.splitext(os.path.basename(options.chunker))[0]
		#print moduleName
		module = imp.load_source(moduleName, options.chunker)
		#print module, dir(module), module.__builtins__, module.__name__
		chunker = module.getChunker(args)
	else:
		chunker = DefaultChunker(1000000)
	compressor = options.compressor.lower()+"compressor"
	zipper = ZipControl(compressor, chunker)
	
	if options.mode[0].lower() == 'c':
		print options.stdin
		if options.stdin:
			obj = sys.stdin
		else:
			obj = open(options.infile)
		zipper.compress(obj, options.outfile)
	elif options.mode[0].lower() == 'd':
		zipper.decompress(options.infile, options.outfile)
		
	#obj = open(sys.argv[1])
	#zipper.compress(obj, "test1.randomaccess.gz.raz")
	#zipper.unzip("test1.randomaccess.gz.raz","blah")
	#obj = zipper.open("test1.randomaccess.gz.raz")
	#obj.seek(1000)
	#print "\nRANDOM ACCESS ZIP FILE\n", obj.read(1000)
	#obj2 = open("blah")
	#obj2.seek(1000)
	#print "\n\nNormal\n", obj2.read(1000)
else:
	fileDir = os.path.dirname(__file__)
	if len(fileDir) == 0:
		fileDir = os.getcwd()
	g_Compressors = getCompressors(fileDir)
